package gobvrtp

import (
	"context"
	"fmt"
	"io"
	"log"
	"time"

	"github.com/bluenviron/gortsplib/v4/pkg/rtptime"
	"github.com/pion/rtcp"
	"github.com/pion/rtp"
	"gogs.smarteye.com/servere/gobvrtp/rtpsdp"
	"gogs.smarteye.com/servere/gobvrtp/track"
)

type Session struct {
	ctx     context.Context
	ctxFunc context.CancelCauseFunc

	sessionTime time.Time

	receiveTracks map[uint32]track.Track
	tracks        map[uint32]track.Track
	rtpDecoder    *rtptime.GlobalDecoder // rtp time

	onNewTrack track.OnNewTrackFunc
	onRTP      track.OnReadRTPFunc
	onRTCP     track.OnReadRTCPFunc
	onFrame    track.OnFrameFunc
	onRTPLoss  track.OnLossRTPFunc

	w io.Writer
}

func NewSession() *Session {
	ctx, ctxFunc := context.WithCancelCause(context.TODO())
	return &Session{
		ctx:     ctx,
		ctxFunc: ctxFunc,

		sessionTime: time.Now(),

		receiveTracks: map[uint32]track.Track{},
		tracks:        map[uint32]track.Track{},
		rtpDecoder:    rtptime.NewGlobalDecoder(),

		onRTP: func(t track.Track, packet *rtp.Packet) {
		},
		onRTCP: func(t track.Track, packets []rtcp.Packet) {
		},
		onFrame: func(t track.Track, frame *track.Frame) {
		},
		onRTPLoss: func(t track.Track, sequenceNumber uint16, diff int) {
		},
	}
}

func (c *Session) SetWriter(w io.Writer) {
	c.w = w
}

// AddReceiveTrack 拉流
func (c *Session) AddReceiveTrack(tracks ...track.Track) {
	for _, t := range tracks {
		t.CtlParma().SetDirection(track.DirectionRecvOnly)
		c.AddTrack(t)
	}
}

// AddSendTrack 推流
func (c *Session) AddSendTrack(tracks ...track.Track) {
	for _, t := range tracks {
		t.CtlParma().SetDirection(track.DirectionSendOnly)
		c.AddTrack(t)
	}
}

func (c *Session) AddTrack(tracks ...track.Track) {
	for _, t := range tracks {
		if t.Direction().RecvOnly() {
			c.receiveTracks[t.SSRC()] = t
			return
		}

		c.addTrack(t)
	}
}

func (c *Session) addTrack(t track.Track) {
	p := t.CtlParma()

	p.SetWriteRTPFunc(
		func(t track.Track, packet *rtp.Packet) error {
			return c.WriteRTP(t, packet)
		},
	)

	p.SetWriteRTCPFunc(
		func(t track.Track, packets []rtcp.Packet) error {
			return c.WriteRTCP(t, packets)
		},
	)

	p.SetPacketPTSFunc(func(t track.Track, pkt *rtp.Packet) (time.Duration, bool) {
		return c.PacketPTS(t, pkt)
	})

	p.SetSessionTimeFunc(func() (time.Time, bool) {
		return c.sessionTime, true
	})

	c.tracks[t.SSRC()] = t
}

func (c *Session) WriteRTCP(t track.Track, packets []rtcp.Packet) error {
	// 更加track 修改rtp的ssrc

	data, err := rtcp.Marshal(packets)
	if err != nil {
		return err
	}

	if _, err = c.w.Write(data); err != nil {
		return err
	}

	return nil
}

// WriteRTP 发送rtp数据
func (c *Session) WriteRTP(t track.Track, pkt *rtp.Packet) error {
	// 更加track 修改rtp的ssrc

	data, err := pkt.Marshal()
	if err != nil {
		return err
	}

	if _, err = c.w.Write(data); err != nil {
		return err
	}

	return nil
}

func (c *Session) CreateOffer() ([]byte, error) {
	var tks []track.Track
	for _, t := range c.tracks {
		tks = append(tks, t)
	}
	for _, t := range c.receiveTracks {
		tks = append(tks, t)
	}

	objectSDP := rtpsdp.NewSDPFromTrack(tks...)

	offerSDP, err := objectSDP.Marshal()

	return offerSDP, err
}

func (c *Session) SetRemoteAnswer(answerSDP []byte) error {

	objectAnswerSDP, err := rtpsdp.ParseSDPObject(answerSDP)
	if err != nil {
		return err
	}

	tracks := objectAnswerSDP.GetTracks()
	for _, t := range tracks {
		c.addTrack(t)

		if c.onNewTrack != nil {
			c.onNewTrack(t)
		}
	}

	for _, t := range c.tracks {
		param := t.CtlParma()
		param.SetOnFrameFunc(c.onFrame)
		param.SetOnLossRTPFunc(c.onRTPLoss)
	}

	return nil
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

// SetOnRTPFunc 接收到rtp的回调
func (c *Session) SetOnRTPFunc(fn track.OnReadRTPFunc) *Session {
	c.onRTP = fn
	return c
}

// SetOnNewTrackFunc 接收到新track
func (c *Session) SetOnNewTrackFunc(fn track.OnNewTrackFunc) *Session {
	c.onNewTrack = fn
	return c
}

// SetOnRTCPFunc 接收到rtcp的回调
func (c *Session) SetOnRTCPFunc(fn track.OnReadRTCPFunc) *Session {
	c.onRTCP = fn
	return c
}

func (c *Session) SetOnFrameFunc(fn track.OnFrameFunc) *Session {
	c.onFrame = fn
	return c
}

func (c *Session) SetOnLossRTPFunc(fn track.OnLossRTPFunc) *Session {
	c.onRTPLoss = fn
	return c
}

func (c *Session) PacketPTS(t track.Track, pkt *rtp.Packet) (time.Duration, bool) {
	return c.rtpDecoder.Decode(t, pkt)
}

func (c *Session) PacketNTP(t track.Track, pkt *rtp.Packet) (time.Time, bool) {
	return t.PacketNTP(pkt)
}

func (c *Session) _onRTP(packet *rtp.Packet) {
	t, ok := c.tracks[packet.SSRC]
	if ok {

		c.onRTP(t, packet)
		t.CtlParma().OnReadRTP(t, packet)
	} else {
		log.Printf("not found track. ssrc: %v", packet.SSRC)
	}
}

func (c *Session) _onRTCP(packets []rtcp.Packet) {
	var t track.Track
	if t != nil {
		t.CtlParma().OnReadRTCP(t, packets)
	}

	// 处理rtcp命令
	for _, packet := range packets {
		pkt := packet

		data, _ := pkt.Marshal()
		var h rtcp.Header
		_ = h.Unmarshal(data)

		if h.Type == rtcp.TypeGoodbye {
			if c.ctxFunc != nil {
				c.ctxFunc(fmt.Errorf("goodbye"))
			}
		}
	}
}

func (c *Session) OnReadPacket(bytes []byte) {
	rtcpPackets, rtpPacket, err := c.parsePacket(bytes)
	if err != nil {
		log.Printf("parsePacket fail. %v", err)
		return
	}

	if rtpPacket != nil {
		c._onRTP(rtpPacket)
	}

	if len(rtcpPackets) > 0 {
		c._onRTCP(rtcpPackets)
	}
}

// 解析packet
func (c *Session) parsePacket(packetData []byte) ([]rtcp.Packet, *rtp.Packet, error) {
	var objectRtcp rtcp.Header
	if err := objectRtcp.Unmarshal(packetData); err == nil {
		if objectRtcp.Type >= 192 && objectRtcp.Type <= 223 {
			if packets, err := rtcp.Unmarshal(packetData); err != nil {
				return nil, nil, err
			} else {
				return packets, nil, nil
			}
		}
	}

	var objectRtp rtp.Packet
	if err := objectRtp.Unmarshal(packetData); err != nil {
		return nil, nil, err
	} else {
		return nil, &objectRtp, nil
	}
}
